import numpy as np
import torch
from torch import nn

bn1 = nn.BatchNorm2d(3)

torch.manual_seed(1)
din = torch.randn(1, 3, 32, 32)

np.savetxt("din.txt", din.detach().numpy().flatten(), fmt="%f")

# for n, p in bn1.named_parameters():
#     print(n)
#     print(p)

dout = bn1(din)
np.savetxt("dout.txt", dout.detach().numpy().flatten(), fmt="%f")

print(torch.sum(din[0][0]) / (32 *32))
print(torch.sum(din[0][1]) / (32 *32))
print(torch.sum(din[0][2]) / (32 *32))

print(torch.var(din[0][0]))
print(torch.var(din[0][1]))
print(torch.var(din[0][2]))
# print(torch.var(din))
